//! 健康检查模块
//! 
//! 提供后端服务器的健康监控功能

use std::sync::Arc;
use std::time::{Duration, Instant};
use std::collections::HashMap;
use tokio::sync::{RwLock, mpsc};
use tokio::time::{interval, timeout};
use anyhow::{Result, anyhow};
use tracing::{info, warn, error, debug};
use reqwest::Client;

use crate::load_balancer::config::{HealthCheckConfig, HealthCheckMethod};
use crate::load_balancer::algorithms::Backend;

/// 后端健康状态
#[derive(Debug, Clone)]
pub struct BackendHealth {
    /// 后端ID
    pub backend_id: String,
    /// 是否健康
    pub is_healthy: bool,
    /// 最后检查时间
    pub last_check: Instant,
    /// 连续成功次数
    pub consecutive_successes: u32,
    /// 连续失败次数
    pub consecutive_failures: u32,
    /// 响应时间
    pub response_time: Option<Duration>,
    /// 错误信息
    pub error_message: Option<String>,
}

/// 健康检查结果
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
    /// 是否成功
    pub success: bool,
    /// 响应时间
    pub response_time: Duration,
    /// 错误信息
    pub error_message: Option<String>,
}

/// 健康检查器
pub struct HealthChecker {
    /// 配置
    config: HealthCheckConfig,
    /// 后端服务器列表
    backends: Arc<RwLock<Vec<Backend>>>,
    /// 健康状态映射
    health_status: Arc<RwLock<HashMap<String, BackendHealth>>>,
    /// HTTP客户端
    http_client: Client,
    /// 停止信号发送器
    stop_sender: Option<mpsc::Sender<()>>,
    /// 停止信号接收器
    stop_receiver: Arc<RwLock<Option<mpsc::Receiver<()>>>>,
}

impl HealthChecker {
    /// 创建新的健康检查器
    pub fn new(config: HealthCheckConfig, backends: Arc<RwLock<Vec<Backend>>>) -> Self {
        let (stop_sender, stop_receiver) = mpsc::channel(1);
        
        let http_client = Client::builder()
            .timeout(config.timeout)
            .build()
            .unwrap_or_else(|_| Client::new());

        Self {
            config,
            backends,
            health_status: Arc::new(RwLock::new(HashMap::new())),
            http_client,
            stop_sender: Some(stop_sender),
            stop_receiver: Arc::new(RwLock::new(Some(stop_receiver))),
        }
    }

    /// 启动健康检查
    pub async fn start(&self) -> Result<()> {
        if !self.config.enabled {
            info!("健康检查已禁用");
            return Ok(());
        }

        info!("启动健康检查，间隔: {:?}", self.config.interval);

        // 初始化健康状态
        self.initialize_health_status().await;

        // 启动健康检查任务
        let backends = self.backends.clone();
        let health_status = self.health_status.clone();
        let config = self.config.clone();
        let http_client = self.http_client.clone();
        
        let mut stop_receiver = self.stop_receiver.write().await.take()
            .ok_or_else(|| anyhow!("健康检查器已经启动"))?;

        tokio::spawn(async move {
            let mut check_interval = interval(config.interval);
            
            loop {
                tokio::select! {
                    _ = check_interval.tick() => {
                        Self::perform_health_checks(
                            &backends,
                            &health_status,
                            &config,
                            &http_client,
                        ).await;
                    }
                    _ = stop_receiver.recv() => {
                        info!("收到停止信号，健康检查任务退出");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// 停止健康检查
    pub async fn stop(&self) -> Result<()> {
        if let Some(sender) = &self.stop_sender {
            sender.send(()).await.map_err(|e| anyhow!("发送停止信号失败: {}", e))?;
        }
        
        info!("健康检查已停止");
        Ok(())
    }

    /// 初始化健康状态
    async fn initialize_health_status(&self) {
        let backends = self.backends.read().await;
        let mut health_status = self.health_status.write().await;
        
        for backend in backends.iter() {
            health_status.insert(
                backend.config.id.clone(),
                BackendHealth {
                    backend_id: backend.config.id.clone(),
                    is_healthy: true, // 初始假设为健康
                    last_check: Instant::now(),
                    consecutive_successes: 0,
                    consecutive_failures: 0,
                    response_time: None,
                    error_message: None,
                },
            );
        }
    }

    /// 执行健康检查
    async fn perform_health_checks(
        backends: &Arc<RwLock<Vec<Backend>>>,
        health_status: &Arc<RwLock<HashMap<String, BackendHealth>>>,
        config: &HealthCheckConfig,
        http_client: &Client,
    ) {
        let backends = backends.read().await;
        let mut futures = Vec::new();

        for backend in backends.iter() {
            let backend_clone = backend.clone();
            let config_clone = config.clone();
            let http_client_clone = http_client.clone();
            
            let future = tokio::spawn(async move {
                let result = Self::check_backend_health(
                    &backend_clone,
                    &config_clone,
                    &http_client_clone,
                ).await;
                
                (backend_clone.config.id.clone(), result)
            });
            
            futures.push(future);
        }

        // 等待所有健康检查完成
        for future in futures {
            if let Ok((backend_id, check_result)) = future.await {
                Self::update_health_status(health_status, &backend_id, check_result, config).await;
            }
        }
    }

    /// 检查单个后端健康状态
    async fn check_backend_health(
        backend: &Backend,
        config: &HealthCheckConfig,
        http_client: &Client,
    ) -> HealthCheckResult {
        let start_time = Instant::now();

        let result = match config.method {
            HealthCheckMethod::Http => {
                Self::http_health_check(backend, config, http_client).await
            }
            HealthCheckMethod::Tcp => {
                Self::tcp_health_check(backend, config).await
            }
            HealthCheckMethod::Grpc => {
                Self::grpc_health_check(backend, config).await
            }
            HealthCheckMethod::Command(ref cmd) => {
                Self::command_health_check(backend, cmd, config).await
            }
        };

        let response_time = start_time.elapsed();

        match result {
            Ok(()) => HealthCheckResult {
                success: true,
                response_time,
                error_message: None,
            },
            Err(e) => HealthCheckResult {
                success: false,
                response_time,
                error_message: Some(e.to_string()),
            },
        }
    }

    /// HTTP健康检查
    async fn http_health_check(
        backend: &Backend,
        config: &HealthCheckConfig,
        http_client: &Client,
    ) -> Result<()> {
        let url = format!("http://{}{}", backend.config.address, config.path);
        
        let response = timeout(
            config.timeout,
            http_client.get(&url).send()
        ).await??;

        let status_code = response.status().as_u16();
        
        if config.expected_codes.contains(&status_code) {
            debug!("HTTP健康检查成功: {} -> {}", backend.config.id, status_code);
            Ok(())
        } else {
            Err(anyhow!(
                "HTTP健康检查失败: {} -> 期望: {:?}, 实际: {}",
                backend.config.id,
                config.expected_codes,
                status_code
            ))
        }
    }

    /// TCP健康检查
    async fn tcp_health_check(
        backend: &Backend,
        config: &HealthCheckConfig,
    ) -> Result<()> {
        let result = timeout(
            config.timeout,
            tokio::net::TcpStream::connect(&backend.config.address)
        ).await;

        match result {
            Ok(Ok(_)) => {
                debug!("TCP健康检查成功: {}", backend.config.id);
                Ok(())
            }
            Ok(Err(e)) => {
                Err(anyhow!("TCP连接失败: {} -> {}", backend.config.id, e))
            }
            Err(_) => {
                Err(anyhow!("TCP连接超时: {}", backend.config.id))
            }
        }
    }

    /// gRPC健康检查
    async fn grpc_health_check(
        backend: &Backend,
        config: &HealthCheckConfig,
    ) -> Result<()> {
        // TODO: 实现gRPC健康检查
        // 这里应该使用gRPC Health Checking Protocol
        warn!("gRPC健康检查暂未实现，使用TCP检查替代");
        Self::tcp_health_check(backend, config).await
    }

    /// 命令行健康检查
    async fn command_health_check(
        backend: &Backend,
        command: &str,
        config: &HealthCheckConfig,
    ) -> Result<()> {
        let cmd_with_backend = command
            .replace("{address}", &backend.config.address.to_string())
            .replace("{id}", &backend.config.id);

        let result = timeout(
            config.timeout,
            tokio::process::Command::new("sh")
                .arg("-c")
                .arg(&cmd_with_backend)
                .output()
        ).await;

        match result {
            Ok(Ok(output)) => {
                if output.status.success() {
                    debug!("命令健康检查成功: {}", backend.config.id);
                    Ok(())
                } else {
                    let stderr = String::from_utf8_lossy(&output.stderr);
                    Err(anyhow!(
                        "命令健康检查失败: {} -> {}",
                        backend.config.id,
                        stderr
                    ))
                }
            }
            Ok(Err(e)) => {
                Err(anyhow!("命令执行失败: {} -> {}", backend.config.id, e))
            }
            Err(_) => {
                Err(anyhow!("命令执行超时: {}", backend.config.id))
            }
        }
    }

    /// 更新健康状态
    async fn update_health_status(
        health_status: &Arc<RwLock<HashMap<String, BackendHealth>>>,
        backend_id: &str,
        check_result: HealthCheckResult,
        config: &HealthCheckConfig,
    ) {
        let mut status_map = health_status.write().await;
        let current_status = status_map.get_mut(backend_id);

        if let Some(status) = current_status {
            status.last_check = Instant::now();
            status.response_time = Some(check_result.response_time);
            status.error_message = check_result.error_message.clone();

            if check_result.success {
                status.consecutive_successes += 1;
                status.consecutive_failures = 0;

                // 检查是否恢复健康
                if !status.is_healthy && status.consecutive_successes >= config.healthy_threshold {
                    status.is_healthy = true;
                    info!("后端服务器恢复健康: {}", backend_id);
                }
            } else {
                status.consecutive_failures += 1;
                status.consecutive_successes = 0;

                // 检查是否变为不健康
                if status.is_healthy && status.consecutive_failures >= config.unhealthy_threshold {
                    status.is_healthy = false;
                    warn!(
                        "后端服务器标记为不健康: {}, 错误: {:?}",
                        backend_id,
                        check_result.error_message
                    );
                }

                error!(
                    "后端健康检查失败: {}, 连续失败: {}, 错误: {:?}",
                    backend_id,
                    status.consecutive_failures,
                    check_result.error_message
                );
            }
        }
    }

    /// 获取健康状态
    pub async fn get_health_status(&self) -> Vec<BackendHealth> {
        let status_map = self.health_status.read().await;
        status_map.values().cloned().collect()
    }

    /// 获取特定后端的健康状态
    pub async fn get_backend_health(&self, backend_id: &str) -> Option<BackendHealth> {
        let status_map = self.health_status.read().await;
        status_map.get(backend_id).cloned()
    }

    /// 添加后端
    pub async fn add_backend(&self, backend_id: String) {
        let mut status_map = self.health_status.write().await;
        
        status_map.insert(
            backend_id.clone(),
            BackendHealth {
                backend_id: backend_id.clone(),
                is_healthy: true,
                last_check: Instant::now(),
                consecutive_successes: 0,
                consecutive_failures: 0,
                response_time: None,
                error_message: None,
            },
        );
    }

    /// 移除后端
    pub async fn remove_backend(&self, backend_id: &str) {
        let mut status_map = self.health_status.write().await;
        status_map.remove(backend_id);
    }

    /// 强制设置后端健康状态
    pub async fn set_backend_health(&self, backend_id: &str, is_healthy: bool) {
        let mut status_map = self.health_status.write().await;
        
        if let Some(status) = status_map.get_mut(backend_id) {
            status.is_healthy = is_healthy;
            status.last_check = Instant::now();
            
            if is_healthy {
                status.consecutive_successes = self.config.healthy_threshold;
                status.consecutive_failures = 0;
            } else {
                status.consecutive_failures = self.config.unhealthy_threshold;
                status.consecutive_successes = 0;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::load_balancer::config::{BackendConfig, ResourceCapacity};
    use std::net::SocketAddr;

    fn create_test_backend(id: &str) -> Backend {
        let mut config = BackendConfig::default();
        config.id = id.to_string();
        config.address = "127.0.0.1:8080".parse::<SocketAddr>().unwrap();
        config.capacity = ResourceCapacity::default();
        
        Backend::new(config)
    }

    #[tokio::test]
    async fn test_health_checker_creation() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
            create_test_backend("backend2"),
        ]));
        
        let config = HealthCheckConfig::default();
        let health_checker = HealthChecker::new(config, backends);
        
        let health_status = health_checker.get_health_status().await;
        assert_eq!(health_status.len(), 0); // 未初始化前为空
    }

    fn create_test_health_config() -> HealthCheckConfig {
        HealthCheckConfig {
            enabled: true,
            interval: Duration::from_millis(100),
            timeout: Duration::from_millis(50),
            healthy_threshold: 2,
            unhealthy_threshold: 3,
            method: HealthCheckMethod::Http {
                path: "/health".to_string(),
                expected_status: 200,
                expected_body: None,
            },
        }
    }

    #[tokio::test]
    async fn test_health_status_initialization() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
            create_test_backend("backend2"),
        ]));
        
        let config = create_test_health_config();
        let health_checker = HealthChecker::new(config, backends);
        
        health_checker.initialize_health_status().await;
        
        let health_status = health_checker.get_health_status().await;
        assert_eq!(health_status.len(), 2, "初始化后应有2个后端的健康状态");
        assert!(health_status.contains_key("backend1"));
        assert!(health_status.contains_key("backend2"));
        
        // 检查初始状态
        let backend1_health = &health_status["backend1"];
        assert!(backend1_health.is_healthy, "初始状态应为健康");
        assert_eq!(backend1_health.consecutive_successes, 0);
        assert_eq!(backend1_health.consecutive_failures, 0);
    }

    #[tokio::test]
    async fn test_backend_health_update() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let config = create_test_health_config();
        let health_checker = HealthChecker::new(config, backends.clone());
        
        health_checker.initialize_health_status().await;
        
        // 更新为成功状态
        let success_result = HealthCheckResult {
            success: true,
            response_time: Duration::from_millis(10),
            error_message: None,
        };
        health_checker.update_backend_health("backend1", success_result).await;
        
        let health_status = health_checker.get_health_status().await;
        let backend1_health = &health_status["backend1"];
        assert_eq!(backend1_health.consecutive_successes, 1);
        assert_eq!(backend1_health.consecutive_failures, 0);
        assert_eq!(backend1_health.response_time, Some(Duration::from_millis(10)));
        
        // 更新为失败状态
        let failure_result = HealthCheckResult {
            success: false,
            response_time: Duration::from_millis(100),
            error_message: Some("Connection refused".to_string()),
        };
        health_checker.update_backend_health("backend1", failure_result).await;
        
        let health_status = health_checker.get_health_status().await;
        let backend1_health = &health_status["backend1"];
        assert_eq!(backend1_health.consecutive_successes, 0);
        assert_eq!(backend1_health.consecutive_failures, 1);
        assert_eq!(backend1_health.error_message, Some("Connection refused".to_string()));
    }

    #[tokio::test]
    async fn test_health_status_transition() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let mut config = create_test_health_config();
        config.unhealthy_threshold = 2;
        config.healthy_threshold = 2;
        let health_checker = HealthChecker::new(config, backends.clone());
        
        health_checker.initialize_health_status().await;
        
        // 连续失败，应该变为不健康
        let failure_result = HealthCheckResult {
            success: false,
            response_time: Duration::from_millis(100),
            error_message: Some("Timeout".to_string()),
        };
        
        // 第一次失败
        health_checker.update_backend_health("backend1", failure_result.clone()).await;
        let health_status = health_checker.get_health_status().await;
        assert!(health_status["backend1"].is_healthy, "第一次失败时应该还是健康的");
        
        // 第二次失败 - 达到不健康阈值
        health_checker.update_backend_health("backend1", failure_result).await;
        let health_status = health_checker.get_health_status().await;
        assert!(!health_status["backend1"].is_healthy, "达到失败阈值后应该变为不健康");
        
        // 验证后端实际状态也被更新
        let backends = backends.read().await;
        assert!(!backends[0].is_healthy(), "后端的实际健康状态也应该被更新");
    }

    #[tokio::test]
    async fn test_health_recovery() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let mut config = create_test_health_config();
        config.unhealthy_threshold = 1;
        config.healthy_threshold = 2;
        let health_checker = HealthChecker::new(config, backends.clone());
        
        health_checker.initialize_health_status().await;
        
        // 先设置为不健康
        let failure_result = HealthCheckResult {
            success: false,
            response_time: Duration::from_millis(100),
            error_message: Some("Error".to_string()),
        };
        health_checker.update_backend_health("backend1", failure_result).await;
        
        let health_status = health_checker.get_health_status().await;
        assert!(!health_status["backend1"].is_healthy);
        
        // 连续成功恢复健康
        let success_result = HealthCheckResult {
            success: true,
            response_time: Duration::from_millis(20),
            error_message: None,
        };
        
        // 第一次成功
        health_checker.update_backend_health("backend1", success_result.clone()).await;
        let health_status = health_checker.get_health_status().await;
        assert!(!health_status["backend1"].is_healthy, "第一次成功时应该还是不健康的");
        
        // 第二次成功 - 达到健康阈值
        health_checker.update_backend_health("backend1", success_result).await;
        let health_status = health_checker.get_health_status().await;
        assert!(health_status["backend1"].is_healthy, "达到成功阈值后应该恢复健康");
        
        // 验证后端实际状态也被更新
        let backends = backends.read().await;
        assert!(backends[0].is_healthy(), "后端的实际健康状态也应该被更新");
    }

    #[tokio::test]
    async fn test_http_health_check() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let config = create_test_health_config();
        let health_checker = HealthChecker::new(config, backends);
        
        // 由于我们无法启动真实的HTTP服务器，这里主要测试方法调用不会崩溃
        let result = health_checker.perform_http_check("http://127.0.0.1:8080/health").await;
        
        // 预期会失败（连接拒绝），但不应该panic
        assert!(!result.success);
        assert!(result.error_message.is_some());
        assert!(result.response_time > Duration::from_millis(0));
    }

    #[tokio::test]
    async fn test_tcp_health_check() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let mut config = create_test_health_config();
        config.method = HealthCheckMethod::Tcp;
        let health_checker = HealthChecker::new(config, backends);
        
        // 测试TCP连接检查
        let result = health_checker.perform_tcp_check("127.0.0.1:8080").await;
        
        // 预期会失败（连接拒绝），但不应该panic
        assert!(!result.success);
        assert!(result.error_message.is_some());
    }

    #[tokio::test]
    async fn test_disabled_health_checker() {
        let backends = Arc::new(RwLock::new(vec![
            create_test_backend("backend1"),
        ]));
        
        let mut config = create_test_health_config();
        config.enabled = false;
        let health_checker = HealthChecker::new(config, backends);
        
        // 启动应该成功但不执行任何检查
        let start_result = health_checker.start().await;
        assert!(start_result.is_ok());
    }

    #[test]
    fn test_backend_health_creation() {
        let health = BackendHealth {
            backend_id: "test".to_string(),
            is_healthy: true,
            last_check: Instant::now(),
            consecutive_successes: 0,
            consecutive_failures: 0,
            response_time: None,
            error_message: None,
        };
        
        assert_eq!(health.backend_id, "test");
        assert!(health.is_healthy);
        assert_eq!(health.consecutive_successes, 0);
        assert_eq!(health.consecutive_failures, 0);
        assert!(health.response_time.is_none());
        assert!(health.error_message.is_none());
    }

    #[test]
    fn test_health_check_result_creation() {
        let success_result = HealthCheckResult {
            success: true,
            response_time: Duration::from_millis(25),
            error_message: None,
        };
        
        assert!(success_result.success);
        assert_eq!(success_result.response_time, Duration::from_millis(25));
        assert!(success_result.error_message.is_none());
        
        let failure_result = HealthCheckResult {
            success: false,
            response_time: Duration::from_millis(100),
            error_message: Some("Connection timeout".to_string()),
        };
        
        assert!(!failure_result.success);
        assert_eq!(failure_result.response_time, Duration::from_millis(100));
        assert_eq!(failure_result.error_message.as_ref().unwrap(), "Connection timeout");
    }
}